//
// Created by 86180 on 2025/1/22.
//

#ifndef AVLTREE_RBTREE_H
#define AVLTREE_RBTREE_H
#include <iostream>
enum Color{
    Red,Black
};

template<typename Key,typename Value>
struct RBTreeNode{
    Color color;
    Key key;
    Value value;
    RBTreeNode* left;
    RBTreeNode* right;
    RBTreeNode* parent;
    RBTreeNode(Key k,Value v):key(k),value(v),left(nullptr),right(nullptr),parent(nullptr),color(Red){}
};


template<typename Key,typename Value>
class RBTree{
private:
    RBTreeNode<Key,Value> * root;
    void leftRotate(RBTreeNode<Key,Value> * x){
        RBTreeNode<Key,Value> * y = x->right;
        x->right = y->left;
        if(y->left != nullptr) y->left->parent = x;
        //一定要先更新父节点再更新子节点
        if(x->parent != nullptr){
            if(x->parent->left == x){
                x->parent->left = y;
            }else{
                x->parent->right = y;
                y->parent = nullptr;
            }
            y->parent = x->parent;
        }else{
            //根节点
            root = y;
        }
        y->left = x;
        x->parent = y;
    }

    void rightRotate(RBTreeNode<Key,Value> * y){
        RBTreeNode<Key,Value> * x = y->left;
        y->left = x->right;
        if(x->right != nullptr) x->right->parent = y;
        //一定要先更新父节点再更新子节点
        if(y->parent != nullptr){
            if(y->parent->left == y){
                y->parent->left = x;
            }else{
                y->parent->right = x;
            }
            x->parent = y->parent;
        }else{
            //根节点
            root = x;
            x->parent = nullptr;
        }
        x->right = y;
        y->parent = x;
    }

    void insertFixUp(RBTreeNode<Key, Value>* z) {
        while (z->parent != nullptr && z->parent->color == Red) {
            if (z->parent == z->parent->parent->left) {
                RBTreeNode<Key, Value>* y = z->parent->parent->right; // 叔叔节点
                if (y != nullptr && y->color == Red) {
                    // 情况1
                    z->parent->color = Black;
                    y->color = Black;
                    z->parent->parent->color = Red;
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->right) {
                        // 情况2
                        z = z->parent;
                        leftRotate(z);
                    }
                    // 情况3
                    z->parent->color = Black;
                    z->parent->parent->color = Red;
                    rightRotate(z->parent->parent);
                }
            } else {
                // 父节点是右子节点，情况对称
                RBTreeNode<Key, Value>* y = z->parent->parent->left; // 叔叔节点
                if (y != nullptr && y->color == Red) {
                    // 情况1
                    z->parent->color = Black;
                    y->color = Black;
                    z->parent->parent->color = Red;
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->left) {
                        // 情况2
                        z = z->parent;
                        rightRotate(z);
                    }
                    // 情况3
                    z->parent->color = Black;
                    z->parent->parent->color = Red;
                    leftRotate(z->parent->parent);
                }
            }
        }
        root->color = Black;
    }
    
    void transplant(RBTreeNode<Key, Value>* u, RBTreeNode<Key, Value>* v) {
        if (u->parent == nullptr) {
            root = v;
        } else if (u == u->parent->left) {
            u->parent->left = v;
        } else {
            u->parent->right = v;
        }
        if (v != nullptr) {
            v->parent = u->parent;
        }
    }

    RBTreeNode<Key, Value>* minimum(RBTreeNode<Key, Value>* node) {
        if (node == nullptr) return nullptr;
        while (node->left != nullptr) {
            node = node->left;
        }
        return node;
    }

    void deleteFixup(RBTreeNode<Key, Value>* x) {
        while (x != root && (x == nullptr || x->color == Black)) {
            if (x == nullptr)break;
            if (x == x->parent->left) {
                RBTreeNode<Key, Value>* w = x->parent->right;
                if (w->color == Red) {
                    w->color = Black;
                    x->parent->color = Red;
                    leftRotate(x->parent);
                    w = x->parent->right;
                }
                if ((w->left == nullptr || w->left->color == Black) &&
                    (w->right == nullptr || w->right->color == Black)) {
                    w->color = Red;
                    x = x->parent;
                } else {
                    if (w->right == nullptr || w->right->color == Black) {
                        if (w->left != nullptr) w->left->color = Black;
                        w->color = Red;
                        rightRotate(w);
                        w = x->parent->right;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Black;
                    if (w->right != nullptr) w->right->color = Black;
                    leftRotate(x->parent);
                    x = root;
                }
            } else {
                RBTreeNode<Key, Value>* w = x->parent->left;
                if (w->color == Red) {
                    w->color = Black;
                    x->parent->color = Red;
                    rightRotate(x->parent);
                    w = x->parent->left;
                }
                if ((w->right == nullptr || w->right->color == Black) &&
                    (w->left == nullptr || w->left->color == Black)) {
                    w->color = Red;
                    x = x->parent;
                } else {
                    if (w->left == nullptr || w->left->color == Black) {
                        if (w->right != nullptr) w->right->color = Black;
                        w->color = Red;
                        leftRotate(w);
                        w = x->parent->left;
                    }
                    w->color = x->parent->color;
                    x->parent->color = Black;
                    if (w->left != nullptr) w->left->color = Black;
                    rightRotate(x->parent);
                    x = root;
                }
            }
        }
        if (x != nullptr) x->color = Black;
    }

    void inorderHelper(RBTreeNode<Key, Value>* node) const {
        if (node == nullptr) return;
        inorderHelper(node->left);
        std::cout << node->key << ":"<<node->value<<"  ";
        inorderHelper(node->right);
    }
    void inorderHelper2(RBTreeNode<Key, Value>* node, std::vector<RBTreeNode<Key, Value>*>& result){
        if (node == nullptr) return;
        inorderHelper2(node->left,result);
        result.push_back(node);
        inorderHelper2(node->right,result);
    }
public:
    RBTree():root(nullptr){}

    void insert(const Key& key,const Value& value){
        RBTreeNode<Key,Value> * z = new RBTreeNode<Key,Value>(key,value);
        RBTreeNode<Key,Value> * y = nullptr;
        RBTreeNode<Key,Value> * x = root;

        while(x != nullptr){
            y = x;
            if (z->key < x->key)
                x = x->left;
            else
                x = x->right;
        }
        z->parent = y;
        if(y == nullptr){
            root = z;
        }else if(z->key<y->key){
            y->left = z;
        }else{
            y->right = z;
        }
        insertFixUp(z);
    }

    void remove(const Key& key) {
        RBTreeNode<Key, Value> *z = root;
        while (z != nullptr) {
            if (key < z->key) {
                z = z->left;
            } else if (key > z->key) {
                z = z->right;
            } else {
                break;
            }
        }
        if (z == nullptr) return; // 没有找到要删除的节点

        RBTreeNode<Key, Value> *y = z;
        RBTreeNode<Key, Value> *x;
        Color yOriginalColor = y->color;

        if (z->left == nullptr) {
            x = z->right;
            transplant(z, z->right);
        } else if (z->right == nullptr) {
            x = z->left;
            transplant(z, z->left);
        } else {
            y = minimum(z->right);
            yOriginalColor = y->color;
            x = y->right;
            if (y->parent == z) {
                if (x != nullptr) x->parent = y;
            } else {
                transplant(y, y->right);
                y->right = z->right;
                if (y->right != nullptr) y->right->parent = y;
            }
            transplant(z, y);
            y->left = z->left;
            if (y->left != nullptr) y->left->parent = y;
            y->color = z->color;
        }

        delete z;

        if (yOriginalColor == Black) {
            deleteFixup(x);
        }
    }

    void inorderTraversal() const {
        inorderHelper(root);
        std::cout << std::endl;
    }

    std::vector<RBTreeNode<Key,Value>*> inorder(){
        std::vector<RBTreeNode<Key,Value>*> res;
        inorderHelper2(root,res);
        return res;
    }

    RBTreeNode<Key, Value>* getminimum(){
        return minimum(root);
    }
};


#endif //AVLTREE_RBTREE_H
